from utils import *
import matplotlib.pyplot as plt
import os
import glob
import random

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

class SatVideoTrainSetLoader(Dataset):
    def __init__(self, dataset_dir, patch_size, img_norm_cfg=None):
        super(SatVideoTrainSetLoader, self).__init__()
        self.dataset_name = 'SatVideoIRSDT'
        self.dataset_dir = dataset_dir
        self.patch_size = patch_size
        
        # 读取训练数据列表
        with open(os.path.join(dataset_dir, 'train.txt'), 'r') as f:
            self.train_list = f.read().splitlines()
        
        # 构建完整的训练样本列表
        self.samples = []
        for item in self.train_list:
            # if '_' in item:  # 格式如 000001_01
            #     video_id, frame_id = item.split('_')
            #     self.samples.append((video_id, frame_id))
            # else:  # 格式如 000002 (需要遍历该视频的所有帧)
            video_id = item
            img_dir = os.path.join(dataset_dir, 'train', video_id, 'img')
            if os.path.exists(img_dir):
                frame_files = sorted(glob.glob(os.path.join(img_dir, '*.png')))
                for frame_file in frame_files:
                    frame_name = os.path.splitext(os.path.basename(frame_file))[0]
                    self.samples.append((video_id, frame_name))
        
        if img_norm_cfg == None:
            # 为SatVideoIRSDT设置默认的归一化配置
            self.img_norm_cfg = {'mean': 0.0, 'std': 1.0}
        else:
            self.img_norm_cfg = img_norm_cfg
            
        self.transform = augumentation()

    def __getitem__(self, idx):
        video_id, frame_id = self.samples[idx]
        
        # 构建图像和掩码路径
        img_path = os.path.join(self.dataset_dir, 'train', video_id, 'img', f'{frame_id}.png')
        mask_path = os.path.join(self.dataset_dir, 'train', video_id, 'mask', f'{frame_id}.png')
        
        try:
            # 读取图像和掩码
            img = Image.open(img_path).convert('L')  # 转换为灰度图
            mask = Image.open(mask_path).convert('L')
        except Exception as e:
            print(f"Error loading {img_path} or {mask_path}: {e}")
            # 返回一个默认的样本
            img = Image.new('L', (256, 256), 0)
            mask = Image.new('L', (256, 256), 0)
        
        # 转换为numpy数组并归一化
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0
        
        # 随机裁剪
        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        
        # 数据增强
        img_patch, mask_patch = self.transform(img_patch, mask_patch)
        
        # 添加通道维度
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis, :]
        
        # 转换为tensor
        img_patch = torch.from_numpy(np.ascontiguousarray(img_patch))
        mask_patch = torch.from_numpy(np.ascontiguousarray(mask_patch))
        
        return img_patch, mask_patch

    def __len__(self):
        return len(self.samples)

class SatVideoTestSetLoader(Dataset):
    def __init__(self, dataset_dir, split='val'):
        super(SatVideoTestSetLoader, self).__init__()
        self.dataset_name = 'SatVideoIRSDT'
        self.dataset_dir = dataset_dir
        self.split = split
        
        # 读取测试数据列表
        split_file = f'{split}.txt' if split in ['val', 'test'] else 'test.txt'
        with open(os.path.join(dataset_dir, split_file), 'r') as f:
            self.test_list = f.read().splitlines()
        
        # 构建完整的测试样本列表
        self.samples = []
        for item in self.test_list:
            video_id = item
            img_dir = os.path.join(dataset_dir, split, video_id, 'img')
            if os.path.exists(img_dir):
                frame_files = sorted(glob.glob(os.path.join(img_dir, '*.png')))
                for frame_file in frame_files:
                    frame_name = os.path.splitext(os.path.basename(frame_file))[0]
                    self.samples.append((video_id, frame_name))
        
        # 默认归一化配置
        self.img_norm_cfg = {'mean': 0.0, 'std': 1.0}

    def __getitem__(self, idx):
        video_id, frame_id = self.samples[idx]
        
        # 构建图像和掩码路径
        img_path = os.path.join(self.dataset_dir, self.split, video_id, 'img', f'{frame_id}.png')
        mask_path = os.path.join(self.dataset_dir, self.split, video_id, 'mask', f'{frame_id}.png')
        
        try:
            # 读取图像和掩码
            img = Image.open(img_path).convert('L')  # 转换为灰度图
            mask = Image.open(mask_path).convert('L')
        except Exception as e:
            print(f"Error loading {img_path} or {mask_path}: {e}")
            # 返回一个默认的样本
            img = Image.new('L', (256, 256), 0)
            mask = Image.new('L', (256, 256), 0)
        
        # 获取原始尺寸
        size = img.size  # (width, height)
        
        # 转换为numpy数组并归一化
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0
        
        # 添加通道维度
        img = img[np.newaxis, :]
        mask = mask[np.newaxis, :]
        
        # 转换为tensor
        img = torch.from_numpy(np.ascontiguousarray(img))
        mask = torch.from_numpy(np.ascontiguousarray(mask))
        
        return img, mask, size, f'{video_id}_{frame_id}'

    def __len__(self):
        return len(self.samples)

class augumentation(object):
    def __call__(self, input, target):
        if random.random() < 0.5:
            input = input[::-1, :].copy()
            target = target[::-1, :].copy()
        if random.random() < 0.5:
            input = input[:, ::-1].copy()
            target = target[:, ::-1].copy()
        if random.random() < 0.5:
            input = input.transpose(1, 0).copy()
            target = target.transpose(1, 0).copy()
        return input, target